package retriever

import (
	"context"
	"fmt"
	"sort"

	"ai-rag/embedding"
	"ai-rag/indexer"

	"github.com/cloudwego/eino/components/model"
	"github.com/cloudwego/eino/schema"
	"github.com/sirupsen/logrus"
)

// Retriever 检索器接口
type Retriever interface {
	Retrieve(ctx context.Context, query string) ([]*schema.Document, error)
}

// VectorRetriever 向量检索器实现
type VectorRetriever struct {
	indexer             indexer.Indexer
	embedder            embedding.Embedder
	topK                int
	similarityThreshold float64
}

// NewVectorRetriever 创建新的向量检索器
func NewVectorRetriever(indexer indexer.Indexer, embedder embedding.Embedder, topK int, similarityThreshold float64) *VectorRetriever {
	return &VectorRetriever{
		indexer:             indexer,
		embedder:            embedder,
		topK:                topK,
		similarityThreshold: similarityThreshold,
	}
}

// Retrieve 检索相关文档
func (r *VectorRetriever) Retrieve(ctx context.Context, query string) ([]*schema.Document, error) {
	logrus.WithField("query", query).WithField("topK", r.topK).Info("开始向量检索")

	// 1. 为查询创建嵌入向量
	queryVector, err := r.embedder.EmbedQuery(ctx, query)
	if err != nil {
		logrus.WithError(err).WithField("query", query).Error("创建查询嵌入向量失败")
		return nil, err
	}

	// 2. 在索引中搜索相似文档
	rawResults, err := r.indexer.Search(ctx, queryVector, r.topK)
	if err != nil {
		logrus.WithError(err).Error("搜索文档失败")
		return nil, err
	}

	// 3. 过滤低于相似度阈值的结果
	var results []*schema.Document
	for _, doc := range rawResults {
		// 获取相似度分数，如果存在的话
		similarity, ok := doc.MetaData["score"].(float64)
		if !ok || similarity >= r.similarityThreshold {
			results = append(results, doc)
		}
	}

	logrus.WithField("query", query).WithField("found", len(results)).Info("向量检索完成")
	return results, nil
}

// MultiQueryRetriever 多查询检索器实现
type MultiQueryRetriever struct {
	vectorRetriever Retriever
	llm             model.BaseChatModel
	maxQueries      int
	rewritePrompt   string
}

// NewMultiQueryRetriever 创建新的多查询检索器
func NewMultiQueryRetriever(vectorRetriever Retriever, llm model.BaseChatModel, maxQueries int, rewritePrompt string) *MultiQueryRetriever {
	return &MultiQueryRetriever{
		vectorRetriever: vectorRetriever,
		llm:             llm,
		maxQueries:      maxQueries,
		rewritePrompt:   rewritePrompt,
	}
}

// Retrieve 检索相关文档
func (r *MultiQueryRetriever) Retrieve(ctx context.Context, query string) ([]*schema.Document, error) {
	logrus.WithField("query", query).WithField("max_queries", r.maxQueries).Info("开始多查询检索")

	// 1. 生成多个查询变体
	queries, err := r.generateQueryVariants(ctx, query)
	if err != nil {
		logrus.WithError(err).Error("生成查询变体失败")
		// 如果生成查询变体失败，回退到原始查询
		return r.vectorRetriever.Retrieve(ctx, query)
	}

	// 2. 对每个查询进行检索
	allResults := make([][]*schema.Document, 0, len(queries))
	for _, q := range queries {
		results, err := r.vectorRetriever.Retrieve(ctx, q)
		if err != nil {
			logrus.WithError(err).WithField("variant_query", q).Error("检索查询变体失败")
			continue
		}
		allResults = append(allResults, results)
	}

	// 3. 融合检索结果
	fusedResults := r.fuseResults(allResults)

	logrus.WithField("query", query).WithField("found", len(fusedResults)).Info("多查询检索完成")
	return fusedResults, nil
}

// generateQueryVariants 生成查询变体
func (r *MultiQueryRetriever) generateQueryVariants(ctx context.Context, originalQuery string) ([]string, error) {
	prompt := fmt.Sprintf(r.rewritePrompt, r.maxQueries)
	prompt += "\n\n原始查询: " + originalQuery
	prompt += "\n\n请列出" + fmt.Sprintf("%d", r.maxQueries) + "个不同的查询变体，每个变体一行:\n1."

	messages := []*schema.Message{
		schema.SystemMessage("你是一个查询改写专家。请基于原始查询生成多个不同的查询表述，这些表述应该涵盖原始查询的不同方面和视角。"),
		schema.UserMessage(prompt),
	}

	response, err := r.llm.Generate(ctx, messages)
	if err != nil {
		return nil, err
	}

	// 解析响应，提取查询变体
	var queries []string
	queries = append(queries, originalQuery) // 始终包含原始查询

	// 简单的解析逻辑，假设响应是按行列出的查询变体
	responseText := response.Content
	lines := splitLines(responseText)

	for _, line := range lines {
		// 去除行号和空格
		line = trimQueryVariant(line)
		if line != "" && !contains(queries, line) {
			queries = append(queries, line)
			if len(queries) >= r.maxQueries {
				break
			}
		}
	}

	return queries, nil
}

// fuseResults 融合多个检索结果
func (r *MultiQueryRetriever) fuseResults(allResults [][]*schema.Document) []*schema.Document {
	// 使用RRF (Reciprocal Rank Fusion) 算法融合结果
	scoreMap := make(map[string]float64)
	docMap := make(map[string]*schema.Document)

	// 为每个文档计算RRF分数
	for _, results := range allResults {
		for rank, doc := range results {
			// 使用文档内容的哈希或唯一标识符作为键
			docID := getDocID(doc)
			docMap[docID] = doc
			// RRF公式: 1 / (k + rank)
			scoreMap[docID] += 1.0 / (60.0 + float64(rank+1))
		}
	}

	// 将文档按分数排序
	type scoredDoc struct {
		doc   *schema.Document
		score float64
	}

	scoredDocs := make([]scoredDoc, 0, len(docMap))
	for docID, score := range scoreMap {
		doc := docMap[docID]
		// 将分数添加到文档元数据
		doc.MetaData["fused_score"] = score
		scoredDocs = append(scoredDocs, scoredDoc{doc: doc, score: score})
	}

	// 按分数降序排序
	sort.Slice(scoredDocs, func(i, j int) bool {
		return scoredDocs[i].score > scoredDocs[j].score
	})

	// 返回排序后的文档
	result := make([]*schema.Document, 0, len(scoredDocs))
	for _, sd := range scoredDocs {
		result = append(result, sd.doc)
	}

	// 限制返回数量为原始topK的1.5倍
	maxResults := r.maxQueries * 3
	if len(result) > maxResults {
		result = result[:maxResults]
	}

	return result
}

// HybridRetriever 混合检索器实现
type HybridRetriever struct {
	retrievers []Retriever
	weights    []float64
	topK       int
}

// NewHybridRetriever 创建新的混合检索器
func NewHybridRetriever(retrievers []Retriever, weights []float64, topK int) *HybridRetriever {
	// 如果没有提供权重，使用均匀权重
	if len(weights) == 0 || len(weights) != len(retrievers) {
		weights = make([]float64, len(retrievers))
		for i := range weights {
			weights[i] = 1.0 / float64(len(retrievers))
		}
	}

	return &HybridRetriever{
		retrievers: retrievers,
		weights:    weights,
		topK:       topK,
	}
}

// Retrieve 检索相关文档
func (r *HybridRetriever) Retrieve(ctx context.Context, query string) ([]*schema.Document, error) {
	logrus.WithField("query", query).WithField("retrievers_count", len(r.retrievers)).Info("开始混合检索")

	// 从多个检索器获取结果
	allResults := make([][]*schema.Document, 0, len(r.retrievers))
	for i, retriever := range r.retrievers {
		results, err := retriever.Retrieve(ctx, query)
		if err != nil {
			logrus.WithError(err).WithField("retriever_index", i).Error("检索器失败")
			continue
		}
		allResults = append(allResults, results)
	}

	// 融合结果
	fusedResults := r.fuseResults(allResults)

	// 限制返回数量
	if len(fusedResults) > r.topK {
		fusedResults = fusedResults[:r.topK]
	}

	logrus.WithField("query", query).WithField("found", len(fusedResults)).Info("混合检索完成")
	return fusedResults, nil
}

// fuseResults 融合多个检索器的结果
func (r *HybridRetriever) fuseResults(allResults [][]*schema.Document) []*schema.Document {
	// 使用加权RRF融合结果
	scoreMap := make(map[string]float64)
	docMap := make(map[string]*schema.Document)

	for i, results := range allResults {
		weight := 1.0
		if i < len(r.weights) {
			weight = r.weights[i]
		}

		for rank, doc := range results {
			docID := getDocID(doc)
			docMap[docID] = doc
			scoreMap[docID] += weight / (60.0 + float64(rank+1))
		}
	}

	// 排序和返回，类似于MultiQueryRetriever中的逻辑
	type scoredDoc struct {
		doc   *schema.Document
		score float64
	}
	scoredDocs := make([]scoredDoc, 0, len(docMap))
	for docID, score := range scoreMap {
		doc := docMap[docID]
		doc.MetaData["hybrid_score"] = score
		scoredDocs = append(scoredDocs, scoredDoc{doc: doc, score: score})
	}

	sort.Slice(scoredDocs, func(i, j int) bool {
		return scoredDocs[i].score > scoredDocs[j].score
	})

	result := make([]*schema.Document, 0, len(scoredDocs))
	for _, sd := range scoredDocs {
		result = append(result, sd.doc)
	}

	return result
}

// 辅助函数
func splitLines(text string) []string {
	var lines []string
	current := ""
	for _, char := range text {
		if char == '\n' {
			if current != "" {
				lines = append(lines, current)
				current = ""
			}
		} else {
			current += string(char)
		}
	}
	if current != "" {
		lines = append(lines, current)
	}
	return lines
}

func trimQueryVariant(line string) string {
	// 去除行号 (如 "1. ")
	for i, char := range line {
		if i > 0 && char == '.' {
			if i+1 < len(line) && line[i+1] == ' ' {
				return line[i+2:]
			}
		}
	}
	// 如果没有行号，直接返回去除首尾空格的行
	return line
}

func contains(list []string, item string) bool {
	for _, v := range list {
		if v == item {
			return true
		}
	}
	return false
}

func getDocID(doc *schema.Document) string {
	// 尝试从元数据获取ID
	if id, ok := doc.MetaData["id"].(string); ok && id != "" {
		return id
	}
	// 如果没有ID，使用内容的前100个字符作为标识
	content := doc.Content
	if len(content) > 100 {
		content = content[:100]
	}
	return content
}
